# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import atexit
import json
import time
import traceback

import absl.app
import absl.flags as flags

from xunkemgmt.support.accounting.logic.query_request import (
    AccountRequest,)
from coin.support.exchange_status.logic.rest_status_checker import (
    check_rest_status,
    ensure_orders_are_cleared,
    ExchangeTestRequest,
    logger,
)
from coin.support.exchange_status.logic.checker_util import \
  convert_to_status_check_request_proto
from coin.proto.coin_market_query_pb2 import (
    ProductTicker,)
from coin.exchange.base.order_enum import (
    is_buy_side_str,
    check_order_side,
)

from coin.support.exchange_status.logic.db_importer import \
  ExchangeStatusImporter

flags.DEFINE_integer('num_trials', 3, 'Number of trials.')

flags.DEFINE_enum('exchange', None, ['Bitmex', 'Okex'], 'Specify exchange name.')

flags.DEFINE_integer('frequency', 60, 'Number of seconds between two consecutive tests.')

flags.DEFINE_integer(
    'duration',
    None,
    'Duration for running checker in seconds. If set to None, the program will '
    'run for ever, otherwise will stop after |duration| seconds.')


def gen_test_order_from_ticker(ticker: ProductTicker, order_side, tick_size, min_order_qty):

  check_order_side(order_side)

  def round_to_ticks(orig_price, prod_tick_size):
    new_price = round(orig_price / prod_tick_size) * tick_size
    return new_price

  if is_buy_side_str(order_side):
    price = ticker.bid * 0.90
  else:
    price = ticker.ask * 1.10

  order = {
      'price': round_to_ticks(price, tick_size),
      'qty': min_order_qty,
      'side': order_side,
  }
  return order


def gen_bitmex_test_orders(ticker: ProductTicker):
  buy_order = gen_test_order_from_ticker(ticker, 'buy', 0.5, 50)
  sell_order = gen_test_order_from_ticker(ticker, 'sell', 0.5, 50)
  return [buy_order, sell_order]


def gen_okex_futures_test_orders(ticker: ProductTicker):
  buy_order = gen_test_order_from_ticker(ticker, 'buy_open', 0.01, 1)
  sell_order = gen_test_order_from_ticker(ticker, 'sell_open', 0.01, 1)
  return [buy_order, sell_order]


def main(argv):
  account_requests = {
      'Bitmex':
          AccountRequest(
              market_type='Futures',
              exchange='Bitmex',
              api_version=None,
              owner='pilot01',
              accounting_currency='BTC',
          ),
      'Okex':
          AccountRequest(
              market_type='Futures',
              exchange='Okex',
              api_version='futures',
              owner='pilot01',
              accounting_currency='BTC',
          ),
  }

  test_request_list = [
      ExchangeTestRequest(
          account_request=account_requests['Bitmex'],
          product='BTC-USD.PERPETUAL',
          gen_test_orders=gen_bitmex_test_orders,
      ),
      ExchangeTestRequest(
          account_request=account_requests['Okex'],
          product='BTC-USD.THIS_WEEK',
          gen_test_orders=gen_okex_futures_test_orders,
      ),
  ]

  for test_request in test_request_list:
    if test_request.account_request.exchange == flags.FLAGS.exchange:
      break
  else:
    raise ValueError('Exchange %s test request not found!' % flags.FLAGS.exchange)

  with open(flags.FLAGS.db_config) as db_config_file:
    db_config = json.load(db_config_file)
  importer = ExchangeStatusImporter.create_connector(db_config)

  atexit.register(ensure_orders_are_cleared, test_request)
  start_time = time.time()
  while True:
    if flags.FLAGS.duration and time.time() - start_time > flags.FLAGS.duration:
      break

    try:
      status_request_proto = convert_to_status_check_request_proto(test_request)
      rest_status_proto = check_rest_status(test_request, flags.FLAGS.num_trials)
      importer.insert_rest_status_history([status_request_proto], [rest_status_proto])
    except Exception as e:
      logger.error('error happened. %s\n %s', e, traceback.format_exc())

    time.sleep(flags.FLAGS.frequency)


if __name__ == '__main__':
  absl.app.run(main)
